import pandas as pd
import numpy as np
import pywt
import matplotlib.pyplot as plt


# ==================== 数据准备（计算排沙量） ====================
def process_data():
    INPUT_FILE = "ori_message.xlsx"
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []

    for year in years:
        df = excel_file.parse(year)
        # 处理列名空格（如'含沙量(kg/m3) '末尾有空格）
        df.columns = df.columns.str.strip()  # 去除列名首尾空格
        required_columns = ['年', '月', '流量(m3/s)', '含沙量(kg/m3)']
        # 检查必要列是否存在
        missing_columns = [col for col in required_columns if col not in df.columns]
        if missing_columns:
            raise ValueError(f"工作表 {year} 缺少必要列：{missing_columns}")

        # 填充年、月列并转换为整数
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)

        # 按年月分组，计算流量和含沙量的月平均值
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3)': 'mean'
        }).reset_index()

        # 计算排沙量（kg/s）= 流量(m³/s) × 含沙量(kg/m³)
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3)']
        all_monthly_data.append(grouped)

    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)

    # 生成日期列（格式：年-月-01）
    result['日期'] = pd.to_datetime(
        result['年'].astype(str) + '-' +
        result['月'].astype(str).str.zfill(2) + '-01'
    )

    return result[['日期', '排沙量(kg/s)']]


# ==================== 主流程 ====================
try:
    df = process_data()
except ValueError as e:
    print(f"数据处理错误: {e}")
    exit(1)

time_series = df['排沙量(kg/s)'].values  # 排沙量时间序列
dates = df['日期'].values  # 时间戳

# ==================== 小波变换参数 ====================
dt = 1  # 时间间隔（月）
scales = np.arange(1, 41)  # 尺度范围（对应周期1-40月）
wavelet = 'morl'  # PyWavelets的Morlet小波

# 执行连续小波变换
[coefficients, _] = pywt.cwt(time_series, scales, wavelet, dt)
coeff_real = np.real(coefficients)  # 提取实部

# ==================== 绘制等值线图 ====================
plt.figure(figsize=(12, 6))
t = np.arange(len(time_series))  # 时间索引
period = scales  # 周期（月）

# 绘制填充等值线（颜色与图15一致）
contourf = plt.contourf(t, period, coeff_real, levels=20, cmap='YlGn', alpha=0.8)
plt.contour(t, period, coeff_real, levels=20, colors='k', linewidths=0.5)

# 时间轴格式化
ax = plt.gca()
ax.xaxis.set_major_locator(plt.MultipleLocator(12))  # 每年一个刻度
formatted_dates = [pd.Timestamp(date).strftime('%Y-%m') for date in dates[::12]]
ax.xaxis.set_major_formatter(plt.FixedFormatter(formatted_dates))

plt.xlabel('时间')
plt.ylabel('周期(月)')
plt.title('2016-2021年每月平均排沙量小波系数实部等值线图')
plt.colorbar(contourf, label='小波系数实部')
plt.ylim(10, 40)  # 锁定纵坐标范围
plt.xlim(-0.5, len(time_series) - 0.5)
plt.tight_layout()
plt.show()